package cn.sc.summer.mybatis.config.datasource;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.sc.summer.mybatis.context.DynamicDataSourceContext;
import cn.sc.summer.mybatis.enums.DataSourceEnum;
import cn.sc.summer.mybatis.properties.DataSourceProperties;
import cn.sc.summer.mybatis.properties.DruidDataSourceProperties;
import cn.sc.summer.mybatis.util.DataSourceUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.xa.DruidXADataSource;
import cn.sc.summer.constant.mybatisplus.datasource.DataSourceConstant;
import cn.sc.summer.exception.enums.CommonExceptionEnum;
import cn.sc.summer.mybatis.service.DataSourceType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.jta.atomikos.AtomikosDataSourceBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 类名：多数据源配置
 *
 * @author a-xin
 * @date 2024/7/30 13:20
 */
@Slf4j
@Configuration
public class DynamicDatasourceConfig {

    @Resource
    private DruidDataSourceProperties druidDataSourceProperties;

    @Value(value = "${seata.enabled:false}")
    private Boolean SEATA_ENABLE;

    @Primary
    @Bean(name = DataSourceConstant.MASTER_DATA_SOURCE)
    public DataSource masterDataSource() {
        log.info("==> Start initialization the master datasource: {}", DataSourceEnum.MASTER.getType());

        Map<String, DataSourceProperties> dataSourcePropertiesMap = druidDataSourceProperties.getDruid();
        CommonExceptionEnum.NORMAL_EXCEPTION.exceptionIf(!dataSourcePropertiesMap.containsKey(DataSourceEnum.MASTER.getType()), "Please set the master datasource!");
        DataSourceProperties dataSourceProperties = dataSourcePropertiesMap.get(DataSourceEnum.MASTER.getType());

        DataSource dataSource;
        if (SEATA_ENABLE) {
            dataSource = DataSourceUtil.createDruidDataSource(dataSourceProperties, new DruidDataSource());
        } else {
            DruidXADataSource druidXADataSource = DataSourceUtil.createDruidDataSource(dataSourceProperties, new DruidXADataSource());
            dataSource = DataSourceUtil.atomikosDataSourceBuild(druidXADataSource, "JTA:" + DataSourceEnum.MASTER.getType());
        }
        checkDataSourceType(dataSource);
        return dataSource;
    }

    /**
     * 多数据源配置
     */
    @Bean(name = DataSourceConstant.DYNAMIC_DATA_SOURCE)
    public DynamicDataSourceContext dynamicDataSource() {
        log.info("==> Start initialization the druid dynamic datasource......");
        Map<Object, Object> dynamic = createDynamicDataSource();
        return new DynamicDataSourceContext(dynamic, masterDataSource());
    }

    @Bean(name = DataSourceConstant.CREATE_DYNAMIC_DATASOURCE)
    public Map<Object, Object> createDynamicDataSource() {
        Map<Object, Object> resultDataSource = new HashMap<>();
        Map<String, DataSourceProperties> druidDataSources = druidDataSourceProperties.getDruid();
        druidDataSources.forEach((dataSourceName, dataSourceProperties) -> {
            log.info("==> Loading the dynamic datasource: {} - [{}]", SEATA_ENABLE ? "SEATA" : "JTA", dataSourceName);
            //这里排除已经初始化的master数据源，否则主数据源会初始化两次
            if (!dataSourceName.equals(DataSourceEnum.MASTER.getType())) {
                if (SEATA_ENABLE) {
                    DruidDataSource druidDataSource = DataSourceUtil.createDruidDataSource(dataSourceProperties, new DruidDataSource());
                    checkDataSourceType(druidDataSource);
                    resultDataSource.put(dataSourceName, druidDataSource);
                } else {
                    DruidXADataSource druidXADataSource = DataSourceUtil.createDruidDataSource(dataSourceProperties, new DruidXADataSource());
                    AtomikosDataSourceBean atomikosDataSourceBean = DataSourceUtil.atomikosDataSourceBuild(druidXADataSource, "JTA:" + dataSourceName);
                    checkDataSourceType(druidXADataSource);
                    resultDataSource.put(dataSourceName, atomikosDataSourceBean);
                }
            }
        });
        return resultDataSource;
    }

    @Bean(name = DataSourceConstant.ALL_DATA_SOURCE_TYPE)
    public List<String> allDataSourceType() {
        List<String> dataSourceTypeList = new ArrayList<>();
        Map<String, DataSourceType> dataSourceTypeMap = SpringUtil.getBeansOfType(DataSourceType.class);
        dataSourceTypeMap.forEach((key, dataSourceType) -> {
            List<String> dataSourceTypes = dataSourceType.getDataSourceType();
            if (CollUtil.isNotEmpty(dataSourceTypes)) {
                dataSourceTypeList.addAll(dataSourceTypes);
            }
        });
        return dataSourceTypeList;
    }

    /**
     * 校验数据源类型是否合法
     *
     * @param dataSource 数据源信息
     */
    private void checkDataSourceType(DataSource dataSource) {
        try {
            List<String> dataSourceTypeList = allDataSourceType();
            if (CollUtil.isEmpty(dataSourceTypeList)) {
                throw new RuntimeException("Specify the type of data source to be accessed! Example implements this class: cn.sc.summer.mybatis.service.DataSourceType");
            }
            Connection connection = dataSource.getConnection();
            String databaseProductName = connection.getMetaData().getDatabaseProductName().toLowerCase();
            if (CharSequenceUtil.isBlank(databaseProductName) || !dataSourceTypeList.contains(databaseProductName)) {
                throw new RuntimeException("Unsupported data source types: " + databaseProductName);
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

}
